clear all
colors={[0 0.4470 0.7410],[0.8500 0.3250 0.0980],[0.9290 0.6940 0.1250],[0.4940 0.1840 0.5560],[0.4660 0.6740 0.1880],[0.3010 0.7450 0.9330],[0.6350 0.0780 0.1840]};


pp = load('mesh_delay_scan_5MHz.mat');
Omega_m = pp.Omega_m;
Delta_m = pp.Delta_m;
mu_m = pp.mu_m;
pg_m = pp.pg_m;
count_m = pp.count_m;

P = [2 1 3];
X = permute(Omega_m,P);
Y = permute(Delta_m,P);
Z = permute(mu_m,P);
Vg = permute(pg_m,P);
Vc = permute(count_m,P);
global F_pg F_count
F_pg = griddedInterpolant(X,Y,Z,Vg,'makima');
F_count = griddedInterpolant(X,Y,Z,Vc,'makima');


% unit of time is us
% unit of energy is MHz 
global v0 v1 v2
v0 = 4.18; %m/s
v1 = 3.45; %m/s
v2 = 0.74; %m/s

% % % initial molecule state is n=1
% v0 = 4.81; %m/s
% v1 = 4.19; %m/s
% v2 = 2.49; %m/s

%% data


% Q2 data: '20230514_100VR_30G_10kHz_REMPI_scan_3'
x2 = [-160  -120   -80   -40   -30   -20   -10     0     5    10    15    20    25    30    40    50];
signal = [7.2500   10.8500   13.0000   22.5500   24.0500   28.2000   32.6000   40.9000   47.5000   50.6000   53.0000   56.1000   50.8000   40.8500   16.4500    3.5500];
e_signal = [0.6500    0.8016    0.8775    1.0897    1.1391    1.2207    1.3077    1.4612    1.5732    1.6155    1.6553    1.6956    1.6294    1.4517    0.9341    0.5025];

bk = [2.5500    1.9000    0.8500    3.8000    3.3000    3.2500    4.0500    4.8000    3.4000    4.2500    4.1000    4.4500    4.7000    4.2000    3.0500    2.5000];
e_bk = [0.4330    0.4416    0.4031    0.5000    0.5099    0.4924    0.5315    0.5745    0.5196    0.5408    0.5431    0.5408    0.5916    0.5244    0.4500    0.4472];
x2 = (-x2+30+20+12-3)*1e-3;

% 
% % Q2 data: '20230621_100VR_30G_10kHz_SR_scan_10'
% signal = [5.7000    9.5500   12.8500   15.9500   18.3000   20.8500   25.4000   30.0500   33.9000   37.7500   39.9500   38.5000   33.9000   24.8500    6.9000    1.7000];
% e_signal = [0.6245    0.7632    0.8352    0.9394    1.0149    1.0642    1.1597    1.2480    1.3266    1.4027    1.4327    1.4089    1.3285    1.1500    0.6403    0.3674];
% bk = [1.0000    1.1500    1.2500    1.2500    0.7500    1.9000    1.1000    1.7000    1.9000    1.3000    2.1000    1.4000    1.7000    1.4500  1.6000    1.8000];
% ebk = [0.3937    0.4031    0.3428    0.3841    0.3905    0.4301    0.3606    0.3742    0.4000    0.3808    0.4000    0.3606    0.3937    0.3905    0.3808    0.3742];
% x2 = (-x2+30+20+12-4)*1e-3;



% % 2G Q2 data: '20230626_100VR_2G_10kHz_SR_scan_1'
% 
% x2 = [-160  -120   -80   -40   -30   -20   -10     0     5    10    15    20    25    30    40    50];
% signal = [8.0000    9.4000   14.1000   18.7000   19.7000   24.3000   24.9000   30.8000   34.1000   31.3000   42.7000   38.8000   40.7000   25.4000   10.0000    2.4000];
% e_signal = [1.0100    1.0677    1.2207    1.4177    1.4595    1.6031    1.6093    1.7833    1.8947    1.8193    2.0761    2.0050    2.0809    1.6125 1.0677    0.6325];
% bk = [0.3000         0    1.4000    0.9000    1.7000    0.7000    1.3000    1.3000    1.1000    0.9000    2.5000    1.6000    1.4000    1.3000 1.0000    1.1000];
% b_bk = [0.5000    0.4472    0.4690    0.4796    0.5745    0.4583    0.4796    0.4796    0.5385    0.5196    0.5385    0.5477    0.6325    0.4359 0.4899    0.5196];



y2 = signal - bk;
ey2 = sqrt(e_signal.^2+e_bk.^2); 

% % 30G 22+ N=1 data: '20230628_100VR_30G_10kHz_SR_scan_10'
% x2= [ -160  -120   -80   -40   -30   -20   -10     0     5    10    15    20    25    30    40    50];
% y2 = [7.0000   10.6667    9.0000   15.3333   18.0000   20.0000   26.3333   33.6667   40.3333   35.5000   49.3333   42.3333   44.0000   30.3333 6.3333    3.0000];
% ey2 = [1.7321    1.8856    1.7951    2.3570    2.4944    2.6247    3.0368    3.4157    3.6667    4.3301    4.0825    3.7859    3.8297    3.1798 1.6667    1.1055];




% x2 = (-x2+30+20+12-3)*1e-3;
% x2 = (-x2+30+20+12-4)*1e-3;

% Q1 data: '20230514_100VR_30G_10kHz_REMPI_scan_2'
x1 = [-160  -120   -80   -40   -30   -20   -10     0     5    10    15    20    25    30    40    50];
signal = [4.0000    8.4000    7.6000   12.9333   13.3333   18.4667   21.0000   28.3333   29.4000   29.0000   33.3333   25.4000   29.0667   20.8667    9.2667    2.1333];
e_signal = [0.6037    0.8055    0.8165    0.9615    1.0022    1.1643    1.2275    1.4095    1.4376    1.4376    1.5261    1.3416    1.4329    1.2166    0.8300    0.4989];
bk = [2.2000    2.3333    2.3333    3.4000    2.8000    2.7333    3.3333    3.6667    4.8000    3.3333    4.1333    3.7333    2.2667    3.7333    2.8667    2.3333];
e_bk = [0.4944    0.4944    0.5617    0.5375    0.5497    0.5538    0.5735    0.5850    0.6532    0.5963    0.6182    0.5963    0.5164    0.5812    0.5121    0.5121];

% % Q1 data '20230622_100VR_30G_10kHz_SR_scan_1'
% x1 = [ -160  -120   -80   -40   -30   -20   -10     0     5    10    15    20    25    30    40    50];
% signal = [4.1111    6.0000    7.8889   11.0000   14.2222   12.1111   20.2222   20.6667   24.4444   22.7778   21.2222   21.0000   19.0000 12.2222    6.2222    2.1111];
% e_signal = [0.7935    0.9686    1.0364    1.1493    1.3699    1.2222    1.5476    1.5556    1.6851    1.6292    1.5986    1.5595    1.5275 1.2472    0.9428    0.6383];
% bk = [0.5556    0.4444    1.4444    1.6667    1.0000    1.5556    0.6667    2.1111    1.8889    3.3333    1.2222    2.1111    2.2222 0.8889    0.6667    1.8889];
% e_bk = [0.4843    0.5666    0.5984    0.5329    0.6383    0.5666    0.4714    0.5984    0.5774    0.7027    0.5774    0.5774    0.6849 0.5443    0.5212    0.6186];

% % 2G Q1 data '20230626_100VR_2G_10kHz_SR_scan_2'
% x1 = [ -160  -120   -80   -40   -30   -20   -10     0     5    10    15    20    25    30    40    50];
% signal = [3.4000    5.0000    6.2000    9.6000   11.8000   14.7000   16.5000   26.6000   27.0000   29.5000   29.0000   24.6000   24.1000   17.2000 5.0000    2.3000];
% e_signal = [0.7211    0.8124    0.9055    1.0954    1.1662    1.2923    1.3379    1.6673    1.6912    1.7578    1.7493    1.6000    1.6031    1.3416 0.8124    0.6708];
% bk = [0.8000    0.6000    1.4000    0.5000    1.9000    1.6000    1.9000    2.5000    2.1000    2.0000    1.4000    1.4000    2.2000    2.6000 0.9000    0.6000];
% e_bk = [0.5099    0.4690    0.5831    0.5385    0.6083    0.6000    0.5745    0.6083    0.6083    0.5831    0.5477    0.4899    0.6164    0.5831 0.5000    0.5292];


y1 = signal - bk;
ey1 = sqrt(e_signal.^2+e_bk.^2);

% % 30G, 22+N=1 data
% x1 = [ -160  -120   -80   -40   -30   -20   -10     0     5    10    15    20    25    30    40    50];
% y1 = [1.8000    5.2000    4.4000    5.2000    6.8000    6.2000    9.6000    8.6000   10.8000   12.6000   11.0000   10.6000   10.4000    9.4000 2.4000    1.6000]
% ey1 = [0.7746    1.1662    1.0198    1.0954    1.2000    1.1832    1.4422    1.3416    1.4967    1.6125    1.5362    1.5362    1.5232    1.4000    0.7483    0.6325];
% 
% % 30G, 22+ N=1, '20230628_100VR_30G_10kHz_SR_scan_33','20230628_100VR_30G_10kHz_SR_scan_11'
% x1 = [ -160  -120   -80   -40   -30   -20   -10     0     5    10    15    20    25    30    40    50];
% y1 = [2.6667    5.1667    3.5833    3.7500    6.0000    6.0833    8.0833    8.3333    9.6667   11.8333   10.2857    8.5000    8.7500    7.9167 2.5000    1.3333];
% ey1 = [0.5528    0.7169    0.5833    0.6614    0.7265    0.7407    0.8620    0.8740    0.9052    1.0138    0.8978    0.8975    0.9091    0.8375  0.5137    0.4082];
% y1 = y1-1;

x1 = (-x1+30+20+12-3)*1e-3;
% x1 = (-x1+30+20+12-6)*1e-3;
% % tic
% xft = [0:0.01:0.1 0.15:0.05:0.25];
% % [c,cerr,fit_val] = fit_wrapper(@(c,t) MC_decay2(c,t),c0,x2',y2',ey2',xft);
% % fit_val = MC_decay2(c,xft);
% % plot(xft,fit_val,'-','linewidth',2,'displayname','fitting');
% c0 = [0.5,6,1];
% % MC_delay2(c0,xft);
% figure(83)
% 
% hold on
% plot(xft,MC_delay2(c0,xft),'-o','linewidth',2,'markersize',8,'displayname','n=1')
% toc

%% fit Q2
if 0
figure(14)
clf
hold on

errorbar(x2,y2,ey2,'o','linewidth',2,'markersize',8,'displayname','n=2')


xft = x2;
% [c,cerr,fit_val] = fit_wrapper(@(c,t) MC_decay2(c,t),c0,x2',y2',ey2',xft);
% fit_val = MC_delay2([0.45 13],xft);
% plot(xft,fit_val,'-','linewidth',2,'displayname','fitting');
options = optimset('Display','iter','MaxIter',60)
fun = @(c) sum((MC_delay2(c,x2)-y2').^2);
c0 = [1.32 10.1361 11.1762]; %[50,8.9,1];% [1.09049   10.3640    1.4792]; %
% c0 = [1. 9.5 12.2];
% c0 = [1.3  9 10];
c2 = fminsearch(fun,c0,options)
fit_val = MC_delay2(c2,xft);
% 
cerr2 = extract_error_bar(@(c,t) MC_delay2(c,t),c2,fit_val,y2',x2)
plot(xft,fit_val,'-o','markersize',8,'linewidth',2,'displayname','fitting');

legend
xlabel('532 center delay (us)')
ylabel('KRb ion counts')
set(gcf,'color','white')
set(gca,'fontsize',16)
end
%% fit Q1
if 1
figure(45)
clf
hold on

errorbar(x1,y1,ey1,'o','linewidth',2,'markersize',8,'displayname','n=1')

xft = x1;
% [c,cerr,fit_val] = fit_wrapper(@(c,t) MC_decay2(c,t),c0,x2',y2',ey2',xft);
% fit_val = MC_delay1([0.4 10 10],xft);
% plot(xft,fit_val,'-','linewidth',2,'displayname','fitting');
options = optimset('Display','iter','MaxIter',100)
fun = @(c) sum((MC_delay1(c,x1)-y1').^2);
c0 = [0.81,10.2,8];
% c0 = [0.8,8.7,10];
c1 = fminsearch(fun,c0,options)


fit_val = MC_delay1(c1,xft);
% fit_val = MC_decay1([0.15 15 3],xft);
cerr1 = extract_error_bar(@(c,t) MC_delay1(c,t),c1,fit_val,y1',x1)
plot(xft,fit_val,'-o','markersize',8,'linewidth',2,'displayname','fitting');

legend
xlabel('pulse length (us)')
ylabel('KRb ion counts')
set(gcf,'color','white')
set(gca,'fontsize',16)
end



%% fit R0
if 0
figure(46)
clf
hold on

errorbar(x0,y0,ey0,'o','linewidth',2,'markersize',8,'displayname','n=0')

xft = x0;
% [c,cerr,fit_val] = fit_wrapper(@(c,t) MC_decay2(c,t),c0,x2',y2',ey2',xft);
% fit_val = MC_decay2(c,xft);
% plot(xft,fit_val,'-','linewidth',2,'displayname','fitting');
options = optimset('Display','iter','MaxIter',100)
fun = @(c) sum((MC_decay0(c,x0)-y0).^2);
c0 = [0.13,10,0.30];
c = fminsearch(fun,c0,options)

% lb = [0.03 8 0.01];
% ub = [0.5 25 1];
% options = optimoptions('fmincon','Display','iter');
% c = fmincon(fun,c0,[],[],[],[],lb,ub)

fit_val = MC_decay0(c,xft);
plot(xft,fit_val,'-o','markersize',8,'linewidth',2,'displayname','fitting');
cerr = extract_error_bar(@(c,t) MC_decay0(c,t),c,fit_val,y0,x0)

legend
xlabel('pulse length (us)')
ylabel('KRb ion counts')
set(gcf,'color','white')
set(gca,'fontsize',16)

end


function y = MC_delay2(c,t)
global v2
    n0 = c(1);
    Omega = c(2);
    delta = c(3);
%     y = n0.*counts_int_angle(t,0.74,Omega,5).*exp(-t.^0.5/tau);
    y = n0.*counts_int_angle(t,v2,Omega,delta);
end

function y = MC_delay1(c,t)
global v1
    n0 = c(1);
    Omega = c(2);
    delta = c(3);
%     y = n0.*counts_int_angle(t,v1,Omega,5).*exp(-t.^0.5/tau);
    y = n0.*counts_int_angle(t,v1,Omega,delta);
end

function y = MC_delay0(c,t)
global v0
    n0 = c(1);
    Omega = c(2);
%     tau = c(3);
    tau2 = c(3);
%     y = n0.*counts_int_angle(t,v0,Omega,5).*exp(-t./tau);
    y = n0.*counts_int_angle0(t,v0,Omega,5,tau2);
end

function ctot = counts_int_angle(mu,v,Rabi0,Delta0)  % measured counts at a certain pulse duration
%     rng(1);
    avg = 1000;
    cor = normrnd(0,1,[avg,3]);
    cor = cor./sqrt(sum(cor.^2,2));
    theta = acos(cor(:,3)); 
    phi = atan(cor(:,2)./cor(:,1));
    t0 = rand(1,avg)*100e-6;
%     r0 = v*t0;
    t_dark = abs(125e-6./(v.*sin(theta).*sin(phi))); %now in unit of s
    
    
    r0 = v*t_dark+v*t0';
    ctot=0;
%     theta = pi/2;s
    for i = 1:length(theta)
        th= theta(i);
        N = round(3.0e-3./sin(th)./(v*100e-6));
        
        ctot = ctot+counts(N,mu,v,th,Rabi0,Delta0,r0(i)*sin(th),t_dark(i));
    end
%     ctot = ctot / avg;
end



function counts = counts(N,mu,v,theta,Rabi0,Delta0,r0,td)  % measured counts at a certain pulse duration
% N is number of pulses.
    n0 = round(td./(100e-6));
%     p_g_0 = exp(-n0.*pulse_duration.^0.5/tau).*((100-pulse_duration)/100).^2;
    p_g_0 = 1;
%     p_g_0 = exp(-1e12*Rabi0^2*2*pi*0.005*(n0*0.2*1e-6)/(15.6e6+4*10*10*1e6/15.6));
    counts = 0;
    p_g = 1;
    rrabi = Rabi(0:N',v*sin(theta),r0);
    [p_e_meta,p_g1_meta] = populations(rrabi,mu,v,theta,Rabi0,Delta0);  
%     size(p_e_meta)
%     size(p_g1_meta)
    for i = 1:N+1
        counts = counts + p_e_meta(:,i).*p_g.*rrabi(i).^2;
%         counts = counts + p_e_meta(:,i).*p_g;
%         counts = counts + p_e_meta(i).*p_g;
        p_g = p_g1_meta(:,i).*p_g;%.*exp(-1e12*Rabi0^2.*rrabi(i).^2*2*pi*0.0001*100e-6/(15.6e6+4*10*10*1e6/15.6)); % remaining populations.
    end
    counts = counts.*p_g_0;
    
end



function [p_ee,p_gg] = populations(Rabi,mu,v,theta,Rabi0,Delta0) 
    % t1 is the delay between 670 REMPI on edge, and t1 is the 670 pulse on time 
    % hard code t1 as 30ns, 

    Omega = Rabi0*Rabi;
    % detuning from Doppler shift
    Delta = v*cos(theta)/3e8*4.49844e8+Delta0;

    ssz = [length(Omega),length(mu)];
    p_ee = Ion_count(repmat(Omega',[1,ssz(2)]),repmat(Delta,ssz),repmat(mu,[ssz(1),1]));
    p_gg = Pg(repmat(Omega',[1,ssz(2)]),repmat(Delta,ssz),repmat(mu,[ssz(1),1]));

end
  
function R = Rabi(n,v,r0)
% n is number of pulses 
    w = 1e-3; % beam waist
    r = r0+n.*v*100e-6;
    I = exp(-2*r.^2./w.^2);
    R = exp(-1*r.^2./w.^2);
end


function [se] = extract_error_bar(func_hand,c_pred,y_pred,y_data,x_data)
    % degrees of freedom in the problem
    dof = length(y_pred) - length(c_pred);

    % standard deviation of the residuals
    sdr = sqrt(sum((y_pred - y_data).^2)/dof);

    % jacobian matrix
    [J,~] = jacobianest(func_hand,c_pred,x_data);

    Sigma = sdr^2*inv(J'*J);

    % Parameter standard errors
    se = sqrt(diag(Sigma))';

%     % which suggest rough confidence intervalues around
%     % the parameters might be...
%     abupper = abfinal + 2*se
%     ablower = abfinal - 2*se
end


function counts1 = Ion_count(Omega,Delta,mu)
    global F_count
    counts1 = F_count(Omega',Delta',mu');
end

function pg = Pg(Omega,Delta,mu)
    global F_pg
    pg = F_pg(Omega',Delta',mu');
% pg = zeros(size(Omega));    
% inds = (Omega*2*pi).^2./(15.6*2*pi+4*(Delta*2*pi).^2./(15.6*2*pi))<=9/15*2*pi;
%     pg(inds) = exp(-length(inds).*(Omega(inds)*2*pi).^2../(15.6*2*pi+4*(Delta(inds)*2*pi).^2./(15.6*2*pi)));
%     pg(~inds) = F_pg(Omega(~inds),Delta(~inds),length(~inds));    
end
